import time
from typing import List, Tuple
import OpenSSL
import josepy as jose
from acme import client
from acme import messages
from acme import challenges
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography import x509
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import hashes

from cryptography.hazmat.primitives import hashes
import logging

logger = logging.getLogger(__name__)

from app.schemas.certificates import CertificateChallenge

# 默认使用 Let's Encrypt Staging 以保证安全，可根据需要调整
# DIRECTORY_URL = "https://acme-staging-v02.api.letsencrypt.org/directory"
DIRECTORY_URL = "https://acme-v02.api.letsencrypt.org/directory"
# 正式环境可改用生产目录：https://acme-v02.api.letsencrypt.org/directory


class FinalizePayload(jose.JSONObjectWithFields):
    """ACME finalize 请求负载，csr 需要 base64url 编码后的 DER。"""
    csr: str = jose.field("csr")

class AcmeClient:
    def __init__(self, private_key_pem: str, email: str, directory_url: str = DIRECTORY_URL):
        self.private_key = serialization.load_pem_private_key(
            private_key_pem.encode("utf-8"), password=None
        )
        self.user_key = jose.JWKRSA(key=self.private_key)
        self.email = email
        self.directory_url = directory_url
        self.net = client.ClientNetwork(self.user_key, user_agent="CertDNSHub/1.0")
        self.directory = client.ClientV2.get_directory(self.directory_url, self.net)
        self.client = client.ClientV2(self.directory, self.net)

    def register_account(self) -> str:
        """注册 ACME 账号并返回注册 URI。"""
        from acme import errors
        
        # 用户要求注册时不要提交邮箱
        reg = messages.NewRegistration.from_data(
            terms_of_service_agreed=True
        )
        
        try:
            regr = self.client.new_account(reg)
        except errors.ConflictError as e:
            # 账号已存在
            account_uri = e.location
            if not account_uri and hasattr(e, 'response'):
                account_uri = e.response.headers.get('Location')
                
            if not account_uri:
                raise ValueError("Account exists but cannot find URI")

            # 构造一个临时的 RegistrationResource 用于查询
            account_body = messages.Registration.from_data(
                key=self.user_key,
                contact=(),
                agreement=None
            )
            account = messages.RegistrationResource(
                body=account_body,
                uri=account_uri
            )
            
            regr = self.client.query_registration(account)
            
        # 关键步骤：把注册结果写回 network client，确保后续请求（如 new_order）
        # 能使用 Key ID（KID）而不是 JWK，符合 ACME v2 的要求。
        self.net.account = regr
        return regr.uri

    def create_order(self, domains: List[str]) -> Tuple[messages.OrderResource, str]:
        """创建新订单，返回 (order, private_key_pem)。"""
        # 生成新的私钥
        key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
        private_key_pem = key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        ).decode("utf-8")
        
        # 构造 CSR
        csr_builder = x509.CertificateSigningRequestBuilder()
        common_name = domains[0]
        csr_builder = csr_builder.subject_name(x509.Name([
            x509.NameAttribute(NameOID.COMMON_NAME, common_name),
        ]))
        san_list = [x509.DNSName(d) for d in domains]
        csr_builder = csr_builder.add_extension(
            x509.SubjectAlternativeName(san_list), critical=False
        )
        csr = csr_builder.sign(key, hashes.SHA256())
        csr_pem = csr.public_bytes(serialization.Encoding.PEM)
        
        # 使用 CSR 创建订单
        order = self.client.new_order(csr_pem)
        return order, private_key_pem

    def _ensure_pollable(self, authz):
        """确保授权对象可用于 client.poll，必要时通过 POST-as-GET 补全内容。"""
        if not authz:
            return None
        if isinstance(authz, messages.AuthorizationResource):
            return authz
        uri = getattr(authz, "uri", None) or authz
        response = self.client.net.post(uri, {})
        authz_body = messages.Authorization.from_json(response.json())
        return messages.AuthorizationResource(body=authz_body, uri=uri)

    def get_challenges(self, order) -> List[CertificateChallenge]:
        challenges_data = []
        # 同时兼容 OrderResource 和 Order body 的 authorizations 字段
        # 如果 order 来自 fetch_order，order.authorizations 可能为空，需要回退到 body.authorizations
        authzs = order.authorizations or order.body.authorizations or []
        
        for authz in authzs:
            # 通过 poll 获取最新状态和挑战
            # poll 返回 (AuthorizationResource, response)
            authz, _ = self.client.poll(self._ensure_pollable(authz))
            domain = authz.body.identifier.value
            
            # 提取 DNS-01 挑战
            dns_challenge = None
            for chall in authz.body.challenges:
                if isinstance(chall.chall, challenges.DNS01):
                    dns_challenge = chall
                    break
            
            if dns_challenge:
                response, validation = dns_challenge.response_and_validation(self.client.net.key)
                # acme 库中的 token 是 bytes，需要转换成 base64url 字符串用于 API 响应
                token_str = jose.b64encode(dns_challenge.chall.token).decode("utf-8")

                challenges_data.append(
                    CertificateChallenge(
                        domain=domain,
                        token=token_str,
                        validation_url=dns_challenge.uri,  # 直接使用 URI 作为验证标识
                        key_authorization=validation,
                        txt_record_name=f"_acme-challenge.{domain}",
                        txt_record_value=validation,
                        authorization_url=getattr(authz, "uri", None),
                    )
                )
        return challenges_data

    def _poll_until_valid_or_ready(self, resource, max_attempts: int = 30, sleep_seconds: int = 2):
        """轮询资源直至状态为 valid 或 ready。"""
        logger.info(f"Polling resource {resource.uri} for status valid/ready")
        for attempt in range(max_attempts):
            # 使用 net.post（POST-as-GET）避免 client.poll 在订单上的兼容性问题
            # （部分 acme 版本会认为 poll 只用于授权并访问 .identifier）
            response = self.client.net.post(resource.uri, {})
            updated_body = resource.body.from_json(response.json())
            updated = type(resource)(body=updated_body, uri=resource.uri)
            
            status = updated.body.status
            logger.debug(f"Poll attempt {attempt+1}/{max_attempts}: status={status}")
            
            if status in (messages.STATUS_VALID, messages.STATUS_READY):
                return updated
            if status == messages.STATUS_INVALID:
                # 不直接打印 body，避免某些版本出现 AttributeError
                error = getattr(updated.body, "error", None)
                logger.error(f"Resource became invalid: {error}")
                raise Exception(f"Resource became invalid. Status: {status}, Error: {error}")
                
            time.sleep(sleep_seconds)
            
        raise TimeoutError(f"Polling timed out after {max_attempts * sleep_seconds} seconds")

    def _wait_for_authorizations_valid(
        self,
        authzs,
        max_attempts: int = 30,
        sleep_seconds: int = 2,
    ) -> None:
        """等待所有授权变为 valid 状态，否则抛出异常。"""
        if not authzs:
            return
        pollable_authzs: List[messages.AuthorizationResource | None] = []
        for authz in authzs:
            pollable_authzs.append(self._ensure_pollable(authz))

        logger.info(f"Waiting for {len(pollable_authzs)} authorizations to become valid")

        for attempt in range(max_attempts):
            all_valid = True
            for idx, authz in enumerate(pollable_authzs):
                if authz is None:
                    continue
                updated_authz, _ = self.client.poll(authz)
                pollable_authzs[idx] = updated_authz
                status = updated_authz.body.status
                
                # Only log if not valid to reduce noise, or log every few attempts
                if attempt % 5 == 0 or status != messages.STATUS_PENDING:
                    logger.info(f"Authz {updated_authz.uri} status: {status}")

                if status == messages.STATUS_INVALID:
                    error = getattr(updated_authz.body, "error", None)
                    error_str = str(error) if error else "Unknown error"
                    logger.error(f"Authorization invalid: {error_str}")
                    raise Exception(f"Authorization failed for {updated_authz.uri}: {error_str}")
                if status != messages.STATUS_VALID:
                    all_valid = False
            
            if all_valid:
                logger.info("All authorizations are valid")
                return
            time.sleep(sleep_seconds)

        raise TimeoutError("等待 ACME 授权进入 valid 状态超时")

    def answer_challenges(self, order) -> None:
        """正确使用签名回答所有 DNS-01 挑战。"""
        authorizations = order.authorizations or order.body.authorizations or []
        logger.info(f"Answering challenges for {len(authorizations)} authorizations")
        
        for authz in authorizations:
            pollable_authz = self._ensure_pollable(authz)
            if pollable_authz is None:
                logger.warning("Found None pollable authz, skipping")
                continue

            authz_res, _ = self.client.poll(pollable_authz)
            status = authz_res.body.status
            logger.info(f"Processing authz {authz_res.uri}, status: {status}")
            
            if status == messages.STATUS_VALID:
                logger.info("Authz already valid, skipping")
                continue
                
            found_dns_challenge = False
            for challb in authz_res.body.challenges:
                # Log challenge type for debugging
                # logger.debug(f"Found challenge type: {type(challb.chall)}")
                
                if isinstance(challb.chall, challenges.DNS01):
                    found_dns_challenge = True
                    response_obj = challb.response(self.client.net.key)
                    logger.info(f"Answering DNS-01 challenge for {authz_res.body.identifier.value}")
                    try:
                        self.client.answer_challenge(challb, response_obj)
                        logger.info(f"Answered challenge: {challb.uri}")
                    except Exception as e:
                        logger.error(f"Failed to answer challenge: {e}")
                        raise
                    break
            
            if not found_dns_challenge and status == messages.STATUS_PENDING:
                 error_msg = f"No DNS-01 challenge found for pending authz {authz_res.uri}"
                 logger.error(error_msg)
                 raise ValueError(error_msg)


    def wait_for_authorizations(
        self, authorization_urls: List[str], max_attempts: int = 15, sleep_seconds: int = 2
    ) -> None:
        """轮询所有授权 URL，直到状态全部变为 valid。"""
        if not authorization_urls:
            return
        self._wait_for_authorizations_valid(
            authorization_urls, max_attempts=max_attempts, sleep_seconds=sleep_seconds
        )

    def fetch_order(self, order_url: str) -> messages.OrderResource:
        """通过 POST-as-GET 获取最新订单信息。"""
        response = self.client.net.post(order_url, {})
        order_body = messages.Order.from_json(response.json())
        return messages.OrderResource(body=order_body, uri=order_url)

    def finalize_order(self, order, csr_pem: bytes) -> messages.OrderResource:
        """
        完整的订单完成流程：
        1. 等待所有授权变为 valid。
        2. 提交 finalize（处理 orderNotReady 等情况）。
        3. 轮询直到订单 valid。
        """
        # 1. 等待授权
        max_retries = 30
        retry_interval = 2
        authzs = order.authorizations or order.body.authorizations or []
        self._wait_for_authorizations_valid(
            authzs,
            max_attempts=max_retries,
            sleep_seconds=retry_interval,
        )

        # 2. 提交 finalize
        # 使用 cryptography 加载 CSR，并构造 ACME 需要的 base64url DER 负载
        csr = x509.load_pem_x509_csr(csr_pem)
        csr_der = csr.public_bytes(serialization.Encoding.DER)
        csr_b64 = jose.b64encode(csr_der).decode("utf-8")
        finalize_payload = FinalizePayload(csr=csr_b64)
        
        # 重新加载订单以获取最新状态和 finalize URL
        response = self.client.net.post(order.uri, {})
        order_body = messages.Order.from_json(response.json())
        orderr = messages.OrderResource(body=order_body, uri=order.uri)

        if orderr.body.status == messages.STATUS_VALID:
            return orderr

        if orderr.body.status != messages.STATUS_READY:
            # 等待 ACME 验证完授权并将订单标记为 ready
            orderr = self._poll_until_valid_or_ready(orderr)
            if orderr.body.status == messages.STATUS_VALID:
                # 轮询期间 ACME 已经签发证书
                return orderr
            if orderr.body.status != messages.STATUS_READY:
                raise Exception(
                    f"Order status {orderr.body.status!r} is not ready for finalization"
                )

        try:
            response = self.client.net.post(orderr.body.finalize, finalize_payload)
        except Exception as e:
            if "orderNotReady" in str(e):
                # 即使前面已检查授权，也在短暂等待后重试一次
                time.sleep(5)
                response = self.client.net.post(orderr.body.finalize, finalize_payload)
            else:
                raise

        # 3. 轮询直到订单 valid
        updated_order_body = messages.Order.from_json(response.json())
        updated_order = messages.OrderResource(body=updated_order_body, uri=orderr.uri)
        
        return self._poll_until_valid_or_ready(updated_order)

    # 为兼容旧调用保留的别名，仍建议尽快更新调用端
    poll_and_finalize = finalize_order

# 默认账号的便捷构造函数
from app.services.datastore import store
import datetime

def get_acme_client() -> AcmeClient:
    account = store.get_default_acme_account()
    return AcmeClient(account["private_key"], account["email"])
